// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <pollux/serializers/compact_row_serializer.h>
#include <pollux/vector/vector_fuzzer.h>

namespace kumo::pollux::test {
namespace {
class RowSerializerBenchmark {
 public:
  void compactRowVectorSerde(
      const RowTypePtr& rowType,
      vector_size_t rangeSize) {
    serializer::CompactRowVectorSerde::registerVectorSerde();
    serialize(rowType, rangeSize);
    deregisterVectorSerde();
  }

 private:
  void serialize(const RowTypePtr& rowType, vector_size_t rangeSize) {
    melon::BenchmarkSuspender suspender;
    auto data = makeData(rowType);
    std::vector<IndexRange> indexRanges;
    for (auto begin = 0; begin < data->size(); begin += rangeSize) {
      indexRanges.push_back(
          IndexRange{begin, std::min(rangeSize, data->size() - begin)});
    }
    suspender.dismiss();

    Scratch scratch;
    auto group = std::make_unique<VectorStreamGroup>(pool_.get(), nullptr);
    group->createStreamTree(rowType, data->size());
    group->append(
        data, melon::Range(indexRanges.data(), indexRanges.size()), scratch);

    std::stringstream stream;
    OStreamOutputStream outputStream(&stream);
    group->flush(&outputStream);
  }

  RowVectorPtr makeData(const RowTypePtr& rowType) {
    VectorFuzzer::Options options;
    options.vectorSize = 1'000;

    const auto seed = 1; // For reproducibility.
    VectorFuzzer fuzzer(options, pool_.get(), seed);

    return fuzzer.fuzzInputRow(rowType);
  }

  std::shared_ptr<memory::MemoryPool> pool_{
      memory::memoryManager()->addLeafPool()};
};

#define VECTOR_SERDE_BENCHMARKS(name, rowType)       \
  BENCHMARK(compact_serialize_1_##name) {            \
    RowSerializerBenchmark benchmark;                \
    benchmark.compactRowVectorSerde(rowType, 1);     \
  }                                                  \
  BENCHMARK(compact_serialize_10_##name) {           \
    RowSerializerBenchmark benchmark;                \
    benchmark.compactRowVectorSerde(rowType, 10);    \
  }                                                  \
  BENCHMARK(compact_serialize_100_##name) {          \
    RowSerializerBenchmark benchmark;                \
    benchmark.compactRowVectorSerde(rowType, 100);   \
  }                                                  \
  BENCHMARK(compact_serialize_1000_##name) {         \
    RowSerializerBenchmark benchmark;                \
    benchmark.compactRowVectorSerde(rowType, 1'000); \
  }

VECTOR_SERDE_BENCHMARKS(
    fixedWidth5,
    ROW({BIGINT(), DOUBLE(), BOOLEAN(), TINYINT(), REAL()}));

VECTOR_SERDE_BENCHMARKS(
    fixedWidth10,
    ROW({
        BIGINT(),
        BIGINT(),
        BIGINT(),
        BIGINT(),
        BIGINT(),
        BIGINT(),
        DOUBLE(),
        BIGINT(),
        BIGINT(),
        BIGINT(),
    }));

VECTOR_SERDE_BENCHMARKS(
    fixedWidth20,
    ROW({
        BIGINT(), BIGINT(), BIGINT(), BIGINT(), BIGINT(), BIGINT(), BIGINT(),
        BIGINT(), BIGINT(), BIGINT(), DOUBLE(), DOUBLE(), DOUBLE(), DOUBLE(),
        DOUBLE(), DOUBLE(), DOUBLE(), DOUBLE(), BIGINT(), BIGINT(),
    }));

VECTOR_SERDE_BENCHMARKS(
    decimal,
    ROW({BIGINT(), DECIMAL(12, 2), DECIMAL(38, 18)}));

VECTOR_SERDE_BENCHMARKS(strings1, ROW({BIGINT(), VARCHAR()}));

VECTOR_SERDE_BENCHMARKS(
    strings5,
    ROW({
        BIGINT(),
        VARCHAR(),
        VARCHAR(),
        VARCHAR(),
        VARCHAR(),
        VARCHAR(),
    }));

VECTOR_SERDE_BENCHMARKS(arrays, ROW({BIGINT(), ARRAY(BIGINT())}));

VECTOR_SERDE_BENCHMARKS(nestedArrays, ROW({BIGINT(), ARRAY(ARRAY(BIGINT()))}));

VECTOR_SERDE_BENCHMARKS(maps, ROW({BIGINT(), MAP(BIGINT(), REAL())}));

VECTOR_SERDE_BENCHMARKS(
    structs,
    ROW({BIGINT(), ROW({BIGINT(), DOUBLE(), BOOLEAN(), TINYINT(), REAL()})}));

} // namespace
} // namespace kumo::pollux::test

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};
  kumo::pollux::memory::MemoryManager::initialize({});
  melon::runBenchmarks();
  return 0;
}
